import wave
import os
import os.path
import shutil
import argparse

def get_args():
    parser = argparse.ArgumentParser(description="""计算ppl值""")
    parser.add_argument('-i','--input', type=str, required=True, default='',
                        help='input txt path')
    parser.add_argument('-audio','--audio_path', type=str, required=True, default='',
                        help='audio path')
    parser.add_argument('-o', '--outdir', type=str, required=True, default='',
                        help='output dir')
    args = parser.parse_args()
    return args

def find_by_pattern(path):
    result_dict = {}
    for root, dirs, files in os.walk(path):
        for f in files:
            #if f.endswith(pattern):
            fpath = os.path.abspath(os.path.join(root, f))
            result_dict[f] = fpath
    return result_dict

#拿到text中的id，与音频文件名称匹配，若匹配则提取音频到另一个文件夹

def get_path(path):
    for root, dirs, files in os.walk(path):
        for f in files:
            fpath = os.path.abspath(os.path.join(root, f))
            print(fpath)
            path = os.path.split(fpath)
            return path[0]  

def process_audio_file(text_file,audio_file,out_file):
    txts_dict = open(text_file, 'r', encoding='utf-8')
    wavs_dict = find_by_pattern(audio_file)
    audio_path = get_path(audio_file)
    for txt_item in txts_dict:
        wav_id = txt_item.split('\t')
        ide = wav_id[0]
        #print(wav_id[0])
        for file_name in wavs_dict:
            #print(file_name)
            if ide == file_name:
                now_path = audio_path + '/'+file_name
                shutil.move(now_path, out_file)


def main():
    args = get_args()
    text_file = args.input
    audio_file = os.path.join(args.audio_path)
    print(audio_file)
    out_file = args.outdir

    process_audio_file(text_file, audio_file, out_file)



if __name__ == "__main__":
    main()
